import glob
import numpy as np
import torch
import os
import cv2
from unet import Unet
import torch.nn as nn
import torch 
from dataPro3 import cellDate
from torchvision.transforms import transforms
import matplotlib.pyplot as plt
import numpy as np 
from PIL import Image
trans = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5])])
target_transform=transforms.ToTensor()
if __name__ == "__main__":
    # 选择设备，有cuda用cuda，没有就用cpu
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # 加载网络，图片单通道，分类为1。
    net = Unet(4, 1)
    # 加载模型参数
    net.load_state_dict(torch.load(r'E:\cellTrainning\Unet_weights_20.pth', map_location='cpu'))
    # 测试模式
    net.eval()
    # 读取所有图片路径
    data_test = r'E:\zhujinchao\uterus\cellTest'
    tran_set=cellDate(data_test, transform = trans,target_transform=target_transform)
    train_loader=torch.utils.data.DataLoader(dataset=tran_set,
                                               batch_size=1, 
                                               shuffle=True)
    net.eval()
    with torch.no_grad():
        i=0
        for image,_ in train_loader:
            i+=1
            pred=net(image).sigmoid()
            pred=torch.squeeze(pred).numpy()
            plt.imshow(pred)
            plt.pause(0.01)
            # plt.savefig(r'E:\zhujinchao\uterus\cellTest\2023-06-25 16.04.45\predmask\{}.png'.format(i))
            plt.show()